{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ca360491",
   "metadata": {},
   "source": [
    "# User Guide on Using FP8\n",
    "\n",
    "JAX supports various FP8 formats, including E4M3 (jnp.float8_e4m3fn) and E5M2\n",
    "(jnp.float8_e5m2). Due to the limited range of FP8 data types, higher-precision\n",
    "data must be scaled to fit within the FP8 representable range, a process known\n",
    "as quantization (Q). Conversely, de-quantization (DQ) rescales the FP8 data back\n",
    "to its original type.\n",
    "\n",
    "While jnp.dot supports FP8 inputs directly, proper quantization and\n",
    "dequantization is needed for optimal performance. Flax provides\n",
    "nn.fp8_ops.Fp8DotGeneral and nn.fp8_ops.Fp8Einsum modules that handle\n",
    "this automatically and can be used with existing layers like nn.Dense.\n",
    "\n",
    "This tutorial will walk you through the basics of how to use it.\n",
    "\n",
    "## Setting up our environment\n",
    "\n",
    "Here, we provide the code necessary to set up the environment for our notebook.\n",
    "Additionally, we define a function to check if the XLA-optimized HLO will indeed\n",
    "call an FP8 dot operation under the hood.\n",
    "\n",
    "*Note: This tutorial relies on the XLA-FP8 feature, which is only supported on\n",
    "NVIDIA Hopper GPUs or later.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "177b91c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import flax\n",
    "import jax\n",
    "import re\n",
    "import pprint\n",
    "from jax import random\n",
    "from jax import numpy as jnp\n",
    "from jax._src import test_util as jtu\n",
    "from flax import linen as nn\n",
    "from flax.linen import fp8_ops\n",
    "\n",
    "e4m3 = jnp.float8_e4m3fn\n",
    "f32 = jnp.float32\n",
    "E4M3_MAX = jnp.finfo(e4m3).max.astype(f32)\n",
    "\n",
    "assert jtu.is_cuda_compute_capability_at_least(\"9.0\")\n",
    "\n",
    "def check_fp8_call(lowered):\n",
    "  hlo = lowered.compile()\n",
    "  if re.search(r\"custom-call\\(f8e4m3fn.*, f8e4m3fn.*\", hlo.as_text()):\n",
    "    print(\"Fp8 call detected!\")\n",
    "  else:\n",
    "    print(\"No Fp8 call!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4adc021f",
   "metadata": {},
   "source": [
    "## FLAX Low Level API\n",
    "\n",
    "The JAX dot operations (e.g. `jnp.dot`) support the FP8 dtype inputs. So it is\n",
    "legal to do the following call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c54c374e",
   "metadata": {},
   "outputs": [],
   "source": [
    "k0, k1 = random.split(random.key(0), 2)\n",
    "a = random.uniform(k0, (16, 32))\n",
    "b = random.uniform(k1, (32, 64))\n",
    "@jax.jit\n",
    "def dot_fp8(a, b):\n",
    "  return jnp.dot(a.astype(e4m3), b.astype(e4m3), preferred_element_type=f32)\n",
    "check_fp8_call(dot_fp8.lower(a, b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "adb22878",
   "metadata": {},
   "source": [
    "However, this approach has two key limitations:\n",
    "\n",
    "1. `jnp.dot` does not support custom scaling factors for operands, defaulting to\n",
    "   a scale of 1.0\n",
    "2. The autodiff does not automatically use E5M2 for gradients and E4M3 for\n",
    "   activations/weights during training, which is the recommended practice\n",
    "\n",
    "To overcome these limitations and implement proper FP8 matrix multiplication, we\n",
    "recommend using the Flax FP8 APIs. Let's start with a basic scaling approach.\n",
    "\n",
    "\n",
    "### Current Scaling\n",
    "\n",
    "Scaling factors are usually defined as `scale = amax(x) / MAX`, where `amax` is\n",
    "an operation to find the absolute maximum value of the tensor, and `MAX` is the\n",
    "maximum value of the representable range of the target dtype. This scaling\n",
    "approach allows us to derive the scaling factors directly from the current\n",
    "operand tensors of the dot product."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0e746e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "@jax.jit\n",
    "def dot_fp8(a, b):\n",
    "  a_scale = jnp.max(jnp.abs(A)) / E4M3_MAX\n",
    "  b_scale = jnp.max(jnp.abs(B)) / E4M3_MAX\n",
    "  a = fp8_ops.quantize(a, e4m3, a_scale, f32)\n",
    "  b = fp8_ops.quantize(b, e4m3, b_scale, f32)\n",
    "\n",
    "  c = jnp.dot(a, b, preferred_element_type=f32)\n",
    "  c = fp8_ops.dequantize(c, f32, a_scale * b_scale)\n",
    "  return c\n",
    "\n",
    "c = dot_fp8(a, b)\n",
    "check_fp8_call(dot_fp8.lower(a, b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "59aca6fe",
   "metadata": {},
   "source": [
    "As shown in the code, we perform quantization (`fp8_ops.quantize`) on the\n",
    "tensors to get the lower precision operands. The `jnp.dot` processes them and\n",
    "accumulates the output in high precision (i.e., the `preferred_element_type`).\n",
    "After that, we multiply the result by the scaling factors to dequantize back to\n",
    "the original range (`fp8_ops.dequantize`). Note that while this example uses\n",
    "E4M3 for both inputs, it is possible to use different FP8 dtypes like E4M3 and\n",
    "E5M2 for the inputs. The quantization method and the scaling factors can also be\n",
    "customized based on application needs.\n",
    "\n",
    "One major issue with the current scaling method is the performance overhead\n",
    "introduced by computing `a_scale` and `b_scale`, which requires additional\n",
    "loading of the operand tensors. To overcome this issue, we recommend the delayed\n",
    "scaling.\n",
    "\n",
    "### Delayed Scaling\n",
    "\n",
    "In delayed scaling, we use a scaling factor associated with an amax history. The\n",
    "scaling factor remains a scalar, but the amax history is a list that stores amax\n",
    "values from recent steps (e.g., 1024 steps). Both tensors are computed from\n",
    "previous steps and maintained in the model parameters.\n",
    "\n",
    "The quantization and dequantization operations for delayed scaling are provided\n",
    "by `fp8_ops.in_q` and `fp8_ops.out_dq` respectively. `fp8_ops.in_q` handles\n",
    "input quantization and update the amax history and scaling factor, while\n",
    "`fp8_ops.out_dq` performs output dequantization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf466308",
   "metadata": {},
   "outputs": [],
   "source": [
    "a_scale = jnp.array(1.0)\n",
    "b_scale = jnp.array(1.0)\n",
    "a_amax_hist = jnp.zeros((1024,))\n",
    "b_amax_hist = jnp.zeros((1024,))\n",
    "\n",
    "@jax.jit\n",
    "def dot_fp8(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist):\n",
    "  a, a_scale = fp8_ops.in_q(f32, e4m3, a, a_scale, a_amax_hist)\n",
    "  b, b_scale = fp8_ops.in_q(f32, e4m3, b, b_scale, b_amax_hist)\n",
    "  \n",
    "  c = jnp.dot(a, b, preferred_element_type=f32)\n",
    "  c = fp8_ops.out_dq(f32, a_scale, b_scale, c)\n",
    "  return c\n",
    "\n",
    "c = dot_fp8(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist)\n",
    "check_fp8_call(dot_fp8.lower(a, a_scale, a_amax_hist, b, b_scale, b_amax_hist))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3bdc038",
   "metadata": {},
   "source": [
    "In this example, we first prepare three pairs of scaling factors and amax\n",
    "histories, treating them as results computed from previous steps. Then, we apply\n",
    "`fp8_ops.in_q` to the input operands of `jnp.dot`, followed by `fp8_ops.out_dq`\n",
    "to the output of `jnp.dot`.\n",
    "\n",
    "\n",
    "## FLAX High Level API\n",
    "Flax provides high-level operations to seamlessly integrate FP8 quantization\n",
    "into existing layers. Instead of manually handling quantization of the delayed\n",
    "scaling (e.g., the maintanence of the amax history and scaling factors), users\n",
    "can simply use these drop-in replacements:\n",
    "\n",
    "* `fp8_ops.Fp8DotGeneral` for `lax.dot_general` operations\n",
    "* `fp8_ops.Fp8Einsum` for `jnp.einsum` operations \n",
    "\n",
    "These operations automatically handle all FP8-related functionality, including\n",
    "quantization/dequantization, scale factor updates, and FP8 dtype selection for\n",
    "both forward and backward passes.\n",
    "\n",
    "Consider the following example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd8d9dba",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = nn.Dense(features=64, dot_general_cls=fp8_ops.Fp8DotGeneral)\n",
    "params = model.init(k0, A)\n",
    "\n",
    "@jax.jit\n",
    "def train_step(var, a): \n",
    "  c = model.apply(var, a)\n",
    "  return jnp.sum(c)\n",
    "\n",
    "check_fp8_call(train_step.lower(params, A))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba280e79",
   "metadata": {},
   "source": [
    "By setting `dot_general_cls=fp8_ops.Fp8DotGeneral`, we replace the\n",
    "default `lax.dot_general` operation in `nn.Dense` with an FP8-enabled version.\n",
    "The model usage remains similar, but now includes additional parameters for FP8\n",
    "quantization: scaling factors and amax history values. The next section explains\n",
    "how to update these FP8-specific parameters.\n",
    "\n",
    "For models that use `jnp.einsum` operations, such as Mixture of Experts (MoE)\n",
    "layers, users can replace them with `fp8_ops.Fp8Einsum` to enable FP8\n",
    "quantization. Here's an example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "961b4549",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Any\n",
    "class FooModule(nn.Module):\n",
    "  einsum: Any = None\n",
    "  @nn.compact\n",
    "  def __call__(self, a, b):\n",
    "    if self.einsum is not None:\n",
    "      einsum_fn = self.einsum()\n",
    "    elif self.einsum is None:\n",
    "      einsum_fn = jnp.einsum\n",
    "    c = einsum_fn(\"mk,kn->mn\", a, b)\n",
    "    return c\n",
    "\n",
    "model = FooModule(einsum=fp8_ops.Fp8Einsum)\n",
    "params = model.init(k0, a, b)\n",
    "\n",
    "@jax.jit\n",
    "def train_step(var, a, b):\n",
    "  c = model.apply(var, a, b)\n",
    "  return jnp.sum(c)\n",
    "\n",
    "check_fp8_call(train_step.lower(params, a, b))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a83b0851",
   "metadata": {},
   "source": [
    "## Manipulate FP8 params\n",
    "\n",
    "The following sections explain the internal FP8 parameters managed by\n",
    "`fp8_ops.Fp8DotGeneral` and `fp8_ops.Fp8Einsum`. These parameters\n",
    "include scaling factors and amax history values that control the FP8\n",
    "quantization process. While most users don't need to interact with these\n",
    "directly, understanding them can be valuable for advanced optimization and\n",
    "debugging.\n",
    "\n",
    "Let's first examine the data structure of `params`. In the code below, we redact\n",
    "the parameter values and then display the PyTree structure."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873799fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "params_structure = flax.core.unfreeze(params).copy()\n",
    "params_structure = flax.traverse_util.flatten_dict(params_structure, sep='/')\n",
    "for key, value in params_structure.items():\n",
    "    params_structure[key] = '*'\n",
    "params_structure = flax.traverse_util.unflatten_dict(params_structure, sep='/')\n",
    "pprint.pprint(params_structure)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "031894dc",
   "metadata": {},
   "source": [
    "The output is as follows:\n",
    "\n",
    "```plaintext\n",
    "{'_overwrite_with_gradient': {'Fp8Einsum_0': {'input_amax_history': '*',\n",
    "                                              'input_scale': '*',\n",
    "                                              'kernel_amax_history': '*',\n",
    "                                              'kernel_scale': '*',\n",
    "                                              'output_grad_amax_history': '*',\n",
    "                                              'output_grad_scale': '*'}}}\n",
    "```\n",
    "\n",
    "In addition to the expected `params`, there is an additional category called\n",
    "`_overwrite_with_gradient`. This category includes three pairs of `amax_history`\n",
    "and `scale` for the activation, kernel, and dot gradient, respectively.\n",
    "\n",
    "### Update gradient of FP8 params\n",
    "Now, we perform one training step to obtain the gradients and see how to use\n",
    "them to update the parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "593fc35f",
   "metadata": {},
   "outputs": [],
   "source": [
    "step_fn = jax.jit(jax.grad(train_step, (0, 1)))\n",
    "\n",
    "grads = step_fn(params, A)\n",
    "\n",
    "params = flax.core.unfreeze(params)\n",
    "params = flax.traverse_util.flatten_dict(params, sep='/')\n",
    "grads = flax.traverse_util.flatten_dict(grads[0], sep='/')\n",
    "\n",
    "for key, value in params.items():\n",
    "  if key.startswith('params'):\n",
    "    params[key] = value + 0.01 * grads[key]\n",
    "  if key.startswith('_overwrite_with_gradient'):\n",
    "    params[key] = grads[key]\n",
    "\n",
    "params = flax.traverse_util.unflatten_dict(params, sep='/')\n",
    "params = flax.core.freeze(params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1a8e2153",
   "metadata": {},
   "source": [
    "The above code demonstrates how to update both `params` and\n",
    "`_overwrite_with_gradient`. For `params`, we use the formula `new_param =\n",
    "old_param + 0.01 * grads`, where `0.01` is the learning rate (or users can use\n",
    "whatever optimizers from `optax`). For `_overwrite_with_gradient`, we simply use\n",
    "the gradient to overwrite the old values.\n",
    "\n",
    "Note that `flax.training.train_state.TrainState` conveniently supports the\n",
    "category of `_overwrite_with_gradient`, so users do not need to modify their\n",
    "scripts if they don't use custom `TrainState`.\n",
    "\n",
    "## Accumulate gradient of FP8 params\n",
    "When the same parameter is used in a branched manner, the autograd mechanism\n",
    "will add their gradients from these branches. This is common in scenarios like\n",
    "pipeline parallelism, where each microbatch shares the same set of parameters\n",
    "for the minibatch. However, for the `_overwrite_with_gradient` parameters, this\n",
    "accumulation by addition is not meaningful. Instead, we prefer custom\n",
    "accumulation by taking the maximum value.\n",
    "\n",
    "To address this, we introduce a custom dtype `fp8_ops.fp32_max_grad`. The basic\n",
    "usage is demonstrated below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d3a86e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "fmax32 = fp8_ops.fp32_max_grad\n",
    "\n",
    "def reuse_fp8_param(x, y, scale, amax_history):\n",
    "  scale = scale.astype(fmax32)\n",
    "  amax_history = amax_history.astype(fmax32)\n",
    "\n",
    "  x = fp8_ops.in_qdq(f32, e4m3, x, scale, amax_history)\n",
    "  y = fp8_ops.in_qdq(f32, e4m3, y, scale, amax_history)\n",
    "  return x + y\n",
    "\n",
    "reuse_fp8_param_fn = jax.grad(reuse_fp8_param, (0, 1, 2, 3))\n",
    "reuse_fp8_param_fn = jax.jit(reuse_fp8_param_fn)\n",
    "\n",
    "_, _, new_ah, new_sf = reuse_fp8_param_fn(2.0, 3.0, a_scale, a_amax_hist)\n",
    "print(new_ah, new_sf)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2321a9bb",
   "metadata": {},
   "source": [
    "In this example, we first cast the `scale` and `amax_history` to\n",
    "`fp8_ops.fp32_max_grad` and then call `fp8_ops.in_qdq` twice using the same pair\n",
    "of `scale` and `amax_history`. During autograd, their gradients from each branch\n",
    "will be taken as the maximum, giving us the correct results of:\n",
    "\n",
    "```plaintext\n",
    "1.0 [3. 0. 0. ... 0. 0. 0.]\n",
    "```\n",
    "\n",
    "If we do not perform the type casting, we get the following result, meaning the\n",
    "gradients of the two branches are added:\n",
    "\n",
    "```plaintext\n",
    "2.0 [5. 0. 0. ... 0. 0. 0.]\n",
    "```\n",
    "\n",
    "This casting is already included if users choose to use the high-level APIs.\n",
    "\n",
    "## Deprecated APIs\n",
    "Previously, we provided APIs like `fp8_ops.quantize_dequantize` for current\n",
    "scaling and `fp8_ops.[in|out]_qdq` for delayed scaling. These were used with\n",
    "high precision dot operations, leveraging an XLA-FP8 feature that\n",
    "pattern-matched QDQ->dot sequences to Q->fp8_cublas_gemm. The corresponding\n",
    "high-level API was called `fp8_ops.Fp8DotGeneralOp`. However, this pattern\n",
    "matching-based solution proved brittle, as the patterns could be easily broken\n",
    "by other XLA optimizations. We recommend users migrate from these deprecated\n",
    "APIs to the newer ones described above.\n",
    "\n",
    "For migration, users should replace:\n",
    "* `fp8_ops.quantize_dequantize -> jnp.dot` with `fp8_ops.quantize -> jnp.dot ->\n",
    "  fp8_ops.dequantize`\n",
    "* `fp8_ops.in_qdq -> jnp.dot -> fp8_ops.out_qdq` with `fp8_ops.in_q -> jnp.dot\n",
    "  -> fp8_ops.out_dq`\n",
    "* `fp8_ops.Fp8DotGeneralOp` with `fp8_ops.Fp8DotGeneral`\n",
    "\n",
    "Additionally, we provide an einsum variant through `fp8_ops.Fp8Einsum`."
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "formats": "ipynb,md:myst"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
